// 本文件生成log文件，仅能在node.js环境下执行
import * as Ld from './log-def.js';
import * as TimeF from "./time-format.js";
import * as fs from 'fs';

let nodeIdGenerator = 0;
function randomKernelName(kernelType) {
    let nodeTypes = ["Add", "Sub", "Mul", "MatMulV2", "BatchMatMul", "Gelu", "GatherV2"];
    let index = Math.floor(Math.random() * nodeTypes.length);
    ++nodeIdGenerator;
    return `${kernelType}_${nodeTypes[index]}/${nodeIdGenerator}_${nodeIdGenerator}`;
}

/**
 * generate random number in standard normal distribution,
 * using polar form (https://en.wikipedia.org/wiki/Box%E2%80%93Muller_transform)
 * @returns {*}
 */
function randomNd() {
    let s = 0;
    let u = 0;
    let v = 0;
    while (s === 0 || s >= 1) {
        u = Math.random() * 2 - 1;
        v = Math.random() * 2 - 1;
        s = (u * u) + (v * v);
    }
    // let z1 = v * Math.sqrt(-2 * Math.log(s) / s);
    return u * Math.sqrt(-2 * Math.log(s) / s); // z0
}

function randomInRange(min, max, generator = Math) {
    return generator.random() * (max - min) + min;
}

class TwoRangeRandom {
    /**
     * 两个区间的均匀随机数发生器，两个区间内部均匀分配：
     * [0, firstRangeMax]概率：firstRangePercent
     * (firstRangeMax, 1]概率：(1 - 1firstRangePercent)
     * @param firstRangeMax
     * @param firstRangePercent
     */
    constructor(firstRangeMax, firstRangePercent) {
        this.firstRangeMax = firstRangeMax;
        this.firstRangePercent = firstRangePercent;
    }
    random() {
        let x = Math.random();
        let y = Math.random();
        while (this.#needDrop(x, y)) {
            x = Math.random();
            y = Math.random();
        }
        return x;
    }
    #needDrop(x, y) {
        if (x <= this.firstRangeMax) {
            return y > this.firstRangePercent;
        } else {
            return y > (1 - this.firstRangePercent);
        }
    }
}

class SizeFaker {
    constructor(meanSize) {
        this.meanSize = meanSize;
        this.stdDev = 1024 * 1024 * 1024;
    }
    random() {
        let size = -1;
        while (size <= 0) {
            size = this.meanSize + (randomNd() * this.stdDev);
        }
        return Math.floor(size / 512) * 512 + 512;
    }
}

class AddressFaker {
    constructor() {
        this.nextAddress = 0x10000000;
    }
    newAddress(size) {
        let addr = this.nextAddress;
        this.nextAddress += size;
        return addr;
    }
}

class StreamFaker {
    constructor(streamNum) {
        this.streamNum = streamNum;
        this.mainStreamPercent = 0.8; // 80% 的事件发生在流0上
    }
    fakeStream() {
        let r = Math.random();
        if (r < this.mainStreamPercent) {
            return 0;
        }
        r -= this.mainStreamPercent;
        r /=  1 - this.mainStreamPercent;
        r *= (this.streamNum - 1)
        return Math.floor(r + 1)
    }
    fakeAnotherStream(excludeId) {
        let streamId = this.fakeStream();
        while (streamId === excludeId) {
            streamId = this.fakeStream();
        }
        return streamId;
    }
}

function logArgComparer(a, b) {
    if (a.timestamp < b.timestamp)  {
        return -1;
    }
    if (a.timestamp > b.timestamp) {
        return 1;
    }
    return 0;
}

class LogGenerator {
    constructor() {
        this.streamNum = 3;
        // 申请、使用、释放都在一条流上
        this.totallySingleStreamNum = 0;
        this.multiStreamHum = 0;

        // 是否随机产生不相关的日志
        this.genIrrelevantLog = false;

        // 随机性：
        // 1. 所有内存的申请、流转、回收、洄游时间都是随机的。
        //    随机的时间区间：2023-01-01T00:00:00.000.000 -> 2023-01-03T00:00:00.000.000
        // 2. 除了totallySingleStreamNum，一块内存多少条流使用、被哪几条流使用是随机的
        // 3. 内存的大小
        // 4. 50%的内存生命周期小于整体的50%，的内存生命周期长于整体的50%
        // 5. 一共3条流(0,1,2)

        this.sizeF = new SizeFaker(100 * 1024 * 1024);
        this.addressF = new AddressFaker();
        this.lifeTimeF = new TwoRangeRandom(0.5, 0.8);
        this.minTime = Date.parse('2023-10-01T00:00:00.000Z');
        this.maxTime = Date.parse('2023-10-01T01:00:00.000Z');
        this.streamIdsToRtsStream = [];
        for (let i = 0; i < this.streamNum; ++i) {
            this.streamIdsToRtsStream[i] = this.addressF.newAddress(512);
        }
    }

    setTotalSingleStreamNum(num) {
        this.totallySingleStreamNum = num;
        return this;
    }
    setMultiStreamNum(num) {
        this.multiStreamHum = num;
        return this;
    }

    #generateLinkStream(logs) {
        for (let streamId = 0; streamId < this.streamIdsToRtsStream.length; ++streamId) {
            logs.push({
                timestamp: this.minTime++,
                generator: Ld.genLinkStream,
                arg: {
                    streamId: streamId,
                    stream: "0x" + this.streamIdsToRtsStream[streamId].toString(16)
                }
            });
        }
    }

    #generateAlloc(streamId) {
        let size = this.sizeF.random();
        let address = this.addressF.newAddress(size);
        return {
            timestamp: randomInRange(this.minTime, this.maxTime, Math),
            generator: Ld.genAllocMem,
            arg: {
                size: size,
                address: address,
                stream: streamId,
                name: randomKernelName("AllocHbmMem")
            }
        };
    }
    #generateBirthRecycle(streamId, minTime, address) {
        return {
            timestamp: randomInRange(minTime, this.maxTime, this.lifeTimeF),
            generator: Ld.genBirthRecycle,
            arg: {
                address: address,
                stream: streamId
            }
        };
    }
    #generateOneStreamAllocAndFree(streamId) {
        let allocLog = this.#generateAlloc(streamId);
        return [
            allocLog,
            this.#generateBirthRecycle(streamId, allocLog.timestamp + 0.001, allocLog.arg.address)
        ];
    }

    #generateAllocAndFree(logs) {
        for (let i = 0; i < this.totallySingleStreamNum; ++i) {
            logs.push(...(this.#generateOneStreamAllocAndFree(0)));
        }
    }

    #generatePoolExpand(logs) {
        logs.sort(logArgComparer);
        let streamIdsToPool = new Map();
        for (let i = 0; i < this.streamNum; ++i) {
            streamIdsToPool.set(i, {poolTotalSize: 0, poolUsedSize: 0});
        }
        let blockAddressToSize = new Map();
        let appendLogs = [];
        for (let logE of logs) {
            if (logE.generator === Ld.genAllocMem) {
                blockAddressToSize.set(logE.arg.address, logE.arg.size);
                let pool = streamIdsToPool.get(logE.arg.stream);
                pool.poolUsedSize += logE.arg.size;
                let expandSize = pool.poolUsedSize - pool.poolTotalSize;
                if (expandSize > 0) {
                    let expandTimestamp = logE.timestamp - 10 - (Math.random() * 100);
                    appendLogs.push({
                        timestamp: expandTimestamp,
                        generator: Ld.genPoolExpand,
                        arg: {
                            stream: "0x" + this.streamIdsToRtsStream[logE.arg.stream].toString(16),
                            size: expandSize,
                            address: this.addressF.newAddress(expandSize)
                        }
                    });
                    pool.poolTotalSize += expandSize;
                }
            } else if(logE.generator === Ld.genBirthRecycle) {
                let pool = streamIdsToPool.get(logE.arg.stream);
                pool.poolUsedSize -= blockAddressToSize.get(logE.arg.address) ?? 0;
            }
        }
        logs.push(...appendLogs);
        logs.sort(logArgComparer);
    }

    #generateMultiStreamMem(logs) {
        let streamF = new StreamFaker(this.streamNum);
        let eventId = 0;
        // generate birth, first-recycle
        for (let i = 0; i < this.multiStreamHum; ++i) {
            let birthStream = streamF.fakeStream();
            let firstRecycleStream = streamF.fakeStream();
            let accessStream = firstRecycleStream;
            if (birthStream === firstRecycleStream) {
                accessStream = streamF.fakeAnotherStream(birthStream);
            }

            let namedLogs = {};
            namedLogs.allocLog = this.#generateAlloc(birthStream);
            namedLogs.birthRecycle = this.#generateBirthRecycle(birthStream,
                namedLogs.allocLog.timestamp + 0.001, namedLogs.allocLog.arg.address);
            namedLogs.wanderLog = {
                timestamp: randomInRange(namedLogs.allocLog.timestamp + 0.001, namedLogs.birthRecycle.timestamp),
                generator: Ld.genStartWandering,
                arg: {
                    srcStream: birthStream,
                    stream: accessStream,
                    address: namedLogs.allocLog.arg.address
                }
            };
            namedLogs.birthStreamLocalRecycle = {
                timestamp: randomInRange(namedLogs.wanderLog.timestamp + 0.001, namedLogs.birthRecycle.timestamp),
                generator: Ld.genLocalRecycle,
                arg: {
                    stream: birthStream,
                    address: namedLogs.allocLog.arg.address
                }
            };
            /*
            let wanderEventId = eventId++;
            namedLogs.sendWander = {
                timestamp: randomInRange(namedLogs.wanderLog.timestamp, namedLogs.birthRecycle.timestamp),
                generator: Ld.genWanderSendEvent,
                arg: {
                    stream: birthStream,
                    eventId: wanderEventId,
                    address: namedLogs.allocLog.arg.address
                }
            };
            namedLogs.waitWander = {
                timestamp: randomInRange(namedLogs.sendWander.timestamp, namedLogs.birthRecycle.timestamp),
                generator: Ld.genWanderWaitEvent,
                arg: {
                    stream: accessStream,
                    eventId: wanderEventId,
                    address: namedLogs.allocLog.arg.address
                }
            };
            namedLogs.accessStreamLocalRecycle = {
                timestamp: randomInRange(namedLogs.waitWander.timestamp + 0.001, namedLogs.birthRecycle.timestamp),
                generator: Ld.genLocalRecycle,
                arg: {
                    stream: accessStream,
                    address: namedLogs.allocLog.arg.address
                }
            };
             */
            namedLogs.accessStreamLocalRecycle = {
                timestamp: randomInRange(namedLogs.wanderLog.timestamp + 0.001, namedLogs.birthRecycle.timestamp),
                generator: Ld.genLocalRecycle,
                arg: {
                    stream: accessStream,
                    address: namedLogs.allocLog.arg.address
                }
            };

            let recycleEventId = eventId++;
            if (birthStream === firstRecycleStream) {
                namedLogs.sendRecycle = {
                    timestamp: randomInRange(namedLogs.accessStreamLocalRecycle.timestamp, namedLogs.birthRecycle.timestamp),
                    generator: Ld.genRecycleSendEvent,
                    arg: {
                        stream: accessStream,
                        eventId: recycleEventId,
                        address: namedLogs.allocLog.arg.address
                    }
                };
                namedLogs.waitRecycle = {
                    timestamp: randomInRange(namedLogs.sendRecycle.timestamp, namedLogs.birthRecycle.timestamp),
                    generator: Ld.genRecycleWaitEvent,
                    arg: {
                        stream: birthStream,
                        eventId: recycleEventId,
                        address: namedLogs.allocLog.arg.address
                    }
                };
            } else {
                namedLogs.sendRecycle = {
                    timestamp: randomInRange(namedLogs.birthStreamLocalRecycle.timestamp, namedLogs.birthRecycle.timestamp),
                    generator: Ld.genRecycleSendEvent,
                    arg: {
                        stream: birthStream,
                        eventId: recycleEventId,
                        address: namedLogs.allocLog.arg.address
                    }
                };
                namedLogs.waitRecycle = {
                    timestamp: randomInRange(namedLogs.sendRecycle.timestamp, namedLogs.birthRecycle.timestamp),
                    generator: Ld.genRecycleWaitEvent,
                    arg: {
                        stream: accessStream,
                        eventId: recycleEventId,
                        address: namedLogs.allocLog.arg.address
                    }
                };
                namedLogs.borrowedRecycle = {
                    timestamp: randomInRange(
                        Math.max(namedLogs.birthStreamLocalRecycle.timestamp, namedLogs.waitRecycle.timestamp),
                        namedLogs.birthRecycle.timestamp),
                    generator: Ld.genBorrowRecycle,
                    arg: {
                        stream: accessStream,
                        address: namedLogs.allocLog.arg.address
                    }
                };
                let migrateEventId = eventId++;
                namedLogs.sendMigration = {
                    timestamp: randomInRange(namedLogs.borrowedRecycle.timestamp, namedLogs.birthRecycle.timestamp),
                    generator: Ld.genMigrationSendEvent,
                    arg: {
                        stream: accessStream,
                        eventId: migrateEventId,
                        address: namedLogs.allocLog.arg.address
                    }
                };
                namedLogs.waitMigration = {
                    timestamp: randomInRange(namedLogs.sendMigration.timestamp, namedLogs.birthRecycle.timestamp),
                    generator: Ld.genMigrationWaitEvent,
                    arg: {
                        stream: birthStream,
                        eventId: migrateEventId,
                        address: namedLogs.allocLog.arg.address
                    }
                };
            }
            logs.push(...(Object.entries(namedLogs).map(e => e[1])));
        }
    }

    generate(filePath) {
        let logs = [];
        this.#generateLinkStream(logs);
        this.#generateAllocAndFree(logs);
        this.#generateMultiStreamMem(logs);
        this.#generatePoolExpand(logs);

        let outConsole = console;
        if (filePath !== undefined) {
            let oStream = fs.createWriteStream(filePath);
            outConsole = new console.Console(oStream, oStream);
        }

        for (let log of logs.map(e => e.generator(Object.assign(e.arg, TimeF.convertTimestampToLogString(e.timestamp))))) {
            outConsole.log(log);
        }
    }
}

new LogGenerator().setTotalSingleStreamNum(50).setMultiStreamNum(10).generate("/home/sn/Downloads/multi_stream_60.log");